#include <zenopy/PythonClient.hpp>
#include <zenox/subroutine.hpp>
#include <zenox/minilog.hpp>
#include <zenox/unicode.hpp>
#include <iostream>
#include <string>
#include <thread>
#include <future>
#include <atomic>
#include <deque>
#include <cstring>
#include <sstream>
#ifdef _WIN32
#include <winsock2.h> // winsock2.h must be included before windows.h (https://blog.csdn.net/lsw1wsl/article/details/70310193)
#include <windows.h>
#include <afunix.h>
#undef min
#undef max
#undef far
#undef near
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>
#endif
#include <json/json.h>

namespace zenopy {

#ifdef _WIN32
namespace {
struct WSAStaticInit {
    bool started = false;

    WSAStaticInit() {
        WORD sockVersion = MAKEWORD(2, 2);
        WSADATA wsdata{};
        if (WSAStartup(sockVersion, &wsdata) != 0) [[unlikely]] {
            zenox::log_error("failed to startup WSA (winsock2)");
        } else {
            started = true;
        }
    }

    ~WSAStaticInit() {
        if (started) {
            WSACleanup();
        }
    }
};
}
#endif

struct PythonClient::Impl {
    std::deque<std::packaged_task<Json::Value(int)>> tasks;
    std::mutex tasks_mtx;
    std::condition_variable cv;
    std::vector<std::jthread> workers;
    std::atomic_int num_sendrecvs = 0;
    std::atomic_int bytes_sends = 0;
    std::atomic_int bytes_recvs = 0;
};

PythonClient::PythonClient()
: m_impl(std::make_unique<Impl>()) {}

PythonClient::~PythonClient() {
    for (auto &worker: m_impl->workers) worker.request_stop();
    m_impl->cv.notify_all();
    for (auto &worker: m_impl->workers) worker.join();
    zenox::log_critical("python client exiting, sent {} bytes, recv {} bytes, count {}",
                        m_impl->bytes_sends.load(), m_impl->bytes_recvs.load(), m_impl->num_sendrecvs.load());
}

void PythonClient::_close(int sock) {
    if (sock != -1) [[likely]] {
        zenox::log_trace("closing socket for python client");
        _send(sock, "{\"type\":\"quit\"}\n", false);
#ifdef _WIN32
        ::closesocket(sock);
#else
        ::close(sock);
#endif
    }
}

void PythonClient::open(const char *socketFile, size_t nWorkers) {
    if (nWorkers <= 0) nWorkers = 1;
    for (size_t i = 0; i < nWorkers; i++) {
        m_impl->workers.push_back(std::jthread([this, socketFile = zenox::format("{}-{}", socketFile, i)] (std::stop_token stop) {
            zenox::log_trace("opening TCP socket (AF_UNIX) for python client: [{}]", socketFile);
            int sock = ::socket(AF_UNIX, SOCK_STREAM, 0);
            if (sock == -1) [[unlikely]] {
                zenox::fatal_error("Failed to create socket object");
            }

            struct sockaddr_un serverAddress{};
            serverAddress.sun_family = AF_UNIX;
            std::strncpy(serverAddress.sun_path, zenox::u8path(socketFile).string().c_str(), sizeof(serverAddress.sun_path) - 1);
            if (::connect(sock, (struct sockaddr *)&serverAddress, sizeof(serverAddress)) < 0) [[unlikely]] {
                zenox::log_error("Failed to connect socket '{}': {}", socketFile, std::strerror(errno));
            }

            while (true) {
                if (stop.stop_requested()) [[unlikely]] break;
                std::unique_lock lck(m_impl->tasks_mtx);
                zenox::log_trace("waiting for task income");
                m_impl->cv.wait(lck, [this, &stop] { return !m_impl->tasks.empty() || stop.stop_requested(); });
                if (stop.stop_requested()) [[unlikely]] break;
                auto pkt = std::move(m_impl->tasks.front());
                m_impl->tasks.pop_front();
                lck.unlock();
                zenox::log_trace("got one job {}, running (socket {})", pkt.valid(), sock);
                pkt(sock);
            }

            zenox::log_trace("stop requested, exiting worker thread");
            _close(sock);
        }));
    }
}

std::string PythonClient::_send(int sock, std::string const &jsonData, bool needRecv) {
    m_impl->num_sendrecvs.fetch_add(1, std::memory_order_relaxed);
    m_impl->bytes_sends.fetch_add(jsonData.size(), std::memory_order_relaxed);
    zenox::log_trace("sending to python server: [{}]", jsonData);
    /* std::replace(jsonData.begin(), jsonData.end(), '\n', ' '); */
    /* jsonData.erase(std::remove(jsonData.begin(), jsonData.end(), '\t'), jsonData.end()); */
    if (::send(sock, jsonData.data(), jsonData.size(), 0) < 0) [[unlikely]] {
        zenox::log_error("Failed to send {} bytes of data: {}", jsonData.length(), std::strerror(errno));
        return "";
    }
    if (!needRecv) [[unlikely]] {
        return "OK";
    }
    std::string result;
    result.clear();
    std::size_t mtu = 1024; // 1KB
    result.resize(mtu);
    std::size_t base = 0;
/* #ifdef _WIN32 */
/*     DWORD timeout = 200; */
/* #else */
/*     struct timeval timeout; */
/*     timeout.tv_sec = 0; */
/*     timeout.tv_usec = 200 * 1000; */
/* #endif */
    while (true) { // 防粘包循环
        /* ::setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char *)&timeout, sizeof(timeout)); */
        zenox::log_trace("receiving data window from {} to {}", base, result.size());
        auto n = ::recv(sock, result.data() + base, result.size() - base, 0);
        zenox::log_trace("got partial data: [{}]", std::string_view(result.data() + base, n));
        if (n <= 0) [[unlikely]] {
            zenox::log_error("Empty reply from server: {}", std::strerror(errno));
            return "";
        }
        m_impl->bytes_recvs.fetch_add(n, std::memory_order_relaxed);
        if (result[base + n - 1] == '\n') {
            result.resize(base + n);
            zenox::log_trace("receiving done, total {} bytes", result.size());
            break;
        }
        base += n;
        if (mtu < 65536) mtu *= 2; // 64KB (since L2 approx 256KB, while base + mtu = 127KB)
        result.resize(base + mtu);
    }
    zenox::log_trace("received from python server: [{}]", result);
    return result;
}

Json::Value PythonClient::exec(std::string code, Json::Value jsonArg) {
    std::packaged_task<Json::Value(int)> pkt([this, code = std::move(code), jsonArg = std::move(jsonArg)] (int sock) {
        return this->_exec(sock, std::move(code), std::move(jsonArg));
    });
    auto fut = pkt.get_future();
    std::unique_lock lck(m_impl->tasks_mtx);
    m_impl->tasks.push_back(std::move(pkt));
    lck.unlock();
    m_impl->cv.notify_one();
    return fut.get();
}

Json::Value PythonClient::_exec(int sock, std::string code, Json::Value jsonArg) {
    Json::Value root;
    root["type"] = "code";
    root["code"] = std::move(code);
    root["args"] = std::move(jsonArg);
    std::string jsonIn;
    Json::EvenFastWriter writer(jsonIn);
    writer.evenFastWrite(root);
    zenox::log_debug("sending to python server:\n=============\n{}\n=============", jsonIn);
    auto retStr = _send(sock, jsonIn);
    if (retStr.empty()) [[unlikely]] {
        zenox::fatal_error("Failed to interact with the server");
    }
    zenox::log_debug("python server replied:\n=============\n{}\n=============", retStr);
    Json::Value result;
    Json::Reader reader;
    reader.parse(retStr, result, false);
    if (result.get("status", "error").asString() == "error") [[unlikely]] {
        zenox::fatal_error("Server returns error: {}", result.get("message", "no error message").asString());
    }
    auto jsonRes = result.get("result", Json::Value());
    return jsonRes;
}

ZENO_SUBPROGRAM(python_client_test)(int argc, char **argv) {
    (void)argc;
    (void)argv;

    PythonClient cli;
    cli.open("/tmp/zenopy-socket");
    auto ret = cli.exec(R"PYTHON(
_zpyio.output = 'hello, 中文!'
)PYTHON", Json::Value());
    zenox::log_critical("Python server replies: {}", ret.asString());

    return 0;
}

}
